import os
import logging
from PyPDF2 import PdfReader, PdfWriter
from PIL import Image, ImageEnhance
from tkinter import filedialog


class PdfWatermarker:
    """
    用于添加水印
    """
    def __init__(self):
        """
        初始化水印处理器
        """
        self.logger = None
        self.output_folder = None
        self.input_folder = None
        self.watermark_path = None
        self.select_input_folder()
        self.select_watermark()
        self.setup_logging()

    def setup_logging(self):
        """设置日志配置"""
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s - %(levelname)s - %(message)s'
        )
        self.logger = logging.getLogger(__name__)

    def add_watermark(self, input_pdf, output_pdf, transparency):
        """
        为单个PDF文件添加水印
        :param input_pdf: 原PDF文件
        :param output_pdf: 加水印后的PDF文件
        :param transparency: 透明度
        :return:
        """
        try:
            # 打开PDF文件
            pdf_reader = PdfReader(open(input_pdf, "rb"))
            pdf_writer = PdfWriter()

            # 获取水印
            watermark = Image.open(self.watermark_path).convert("RGBA")

            # 条整水印透明度
            alpha = watermark.split()[3]
            alpha = ImageEnhance.Brightness(alpha).enhance(transparency)
            watermark.putalpha(alpha)

            # 处理每一页
            for page_num in range(len(pdf_reader.pages)):
                page = pdf_reader.pages[page_num]
                pdf_page_size = page.mediabox
                pdf_width = int(float(pdf_page_size.width) * 1.3)
                pdf_height = int(float(pdf_page_size.height) * 1.4)

                # 重置水印大小使其适应页面
                watermark_ratio_width = pdf_width / watermark.width
                watermark_ratio_height = pdf_height / watermark.height
                watermark_width = int(watermark.width * watermark_ratio_width)
                watermark_height = int(watermark.height * watermark_ratio_height)
                watermark_resized = watermark.resize((watermark_width, watermark_height), Image.LANCZOS)

                # 在PDF中创建新的图像
                new_image = Image.new("RGBA", (pdf_width, pdf_height), (255, 255, 255, 0))

                # 将水印居中插入
                pos_x = (pdf_width - watermark_width) // 2
                pos_y = (pdf_height - watermark_height) // 2
                new_image.paste(watermark_resized, (pos_x, pos_y), watermark_resized)

                temp_watermark_pdf = "temp_watermark.pdf"
                new_image.save(temp_watermark_pdf, "PDF", resolution=100.0)

                watermark_pdf_reader = PdfReader(open(temp_watermark_pdf, 'rb'))
                watermark_page = watermark_pdf_reader.pages[0]

                page.merge_page(watermark_page)
                pdf_writer.add_page(page)

            with open(output_pdf, 'wb') as out_file:
                pdf_writer.write(out_file)

        except Exception as e:
            self.logger.error(f'添加水印时发生错误: {str(e)}')
            raise

    def process_folder(self):
        """处理文件夹中的所有PDF文件"""
        try:
            files = [f for f in os.listdir(self.input_folder) if f.endswith(".pdf")]
            total_files = len(files)

            if not os.path.exists(self.output_folder):
                os.makedirs(self.output_folder)

            for i, filename in enumerate(files):
                input_pdf_path = os.path.join(self.input_folder, filename)
                output_pdf_path = os.path.join(self.output_folder, filename)
                self.add_watermark(input_pdf_path, output_pdf_path, 0.6)
                print(f'已经为 {filename}添加水印')

        except Exception as e:
            self.logger.error(f'处理文件夹时发生错误: {str(e)}')
            raise

    def select_input_folder(self):
        """
        选择需要处理的PDF文件夹
        :return: None
        """
        self.input_folder = filedialog.askdirectory(title="请选择PDF文件夹")

        # 获取输入文件夹的父目录
        parent_dir = os.path.dirname(self.input_folder)

        # 获取输入文件夹的名称
        folder_name = os.path.basename(self.input_folder)

        # 在同一层级创建 _output 文件夹
        self.output_folder = os.path.join(parent_dir, f"{folder_name}_output")

        # 创建输出文件夹（如果不存在）
        os.makedirs(self.output_folder, exist_ok=True)

    def select_watermark(self):
        """
        选择水印文件
        :return:
        """
        self.watermark_path = filedialog.askopenfilename(title="请选择水印图片")


if __name__ == '__main__':
    # 实例化水印添加工具
    pdfWterMarker = PdfWatermarker()
    # 调用批量添加水印方法
    pdfWterMarker.process_folder()
